from __future__ import annotations
from rasa.nlu.tokenizers.tokenizer import Tokenizer, Token
import unicodedata
from rasa.shared.nlu.training_data.message import Message
from typing import Text, List, Dict, Any
from rasa.engine.recipes.default_recipe import DefaultV1Recipe
from rasa.engine.graph import ExecutionContext
from rasa.engine.storage.resource import Resource
from rasa.engine.storage.storage import ModelStorage
import spacy

class BaseTokenizer:
    """适应中文的tokenizer
    - 中文基于字
    - 英文基于词
    - 数字分开
    """
    def __init__(self,
                do_lower_case: bool = True,
                return_offset_mapping: bool = True,
                never_split: List = [],
                **kwargs):
        self.do_lower_case = do_lower_case
        self.never_split = never_split
        self.return_offset_mapping = return_offset_mapping


    def _whitespace_tokenize(self, text):
        """去除文本中的空白符"""
        text = text.strip()
        if not text:
            return []
        tokens = text.split()
        return tokens


    def _run_strip_accents(self, text):
        """Strips accents from a piece of text."""
        text = unicodedata.normalize("NFD", text)
        output = []
        for char in text:
            cat = unicodedata.category(char)
            if cat == "Mn":
                continue
            output.append(char)
        return "".join(output)

    def _run_split_on_punc(self, text):
        """Splits punctuation on a piece of text."""
        if text in self.never_split:
            return [text]
        chars = list(text)
        i = 0
        start_new_word = True
        output = []
        while i < len(chars):
            char = chars[i]
            if self._is_punctuation(char):
                output.append([char])
                start_new_word = True
            else:
                if start_new_word:
                    output.append([])
                start_new_word = False
                output[-1].append(char)
            i += 1
        return ["".join(x) for x in output]

    def _tokenize_chinese_chars(self, text):
        """Adds whitespace around any CJK character."""
        output = []
        for char in text:
            cp = ord(char)
            if self._is_chinese_char(cp) or char.isdigit():
                output.append(" ")
                output.append(char)
                output.append(" ")
            else:
                output.append(char)
        return "".join(output)


    def _clean_text(self, text):
        """Performs invalid character removal and whitespace cleanup on text."""
        output = []
        for char in text:
            cp = ord(char)
            if cp == 0 or cp == 0xfffd or self._is_control(char):
                continue
            if self._is_whitespace(char):
                output.append(" ")
            else:
                output.append(char)
        return "".join(output)


    def _is_control(self, char):
        """Checks whether `chars` is a control character."""
        # These are technically control characters but we count them as whitespace
        # characters.
        if char == "\t" or char == "\n" or char == "\r":
            return False
        cat = unicodedata.category(char)
        if cat.startswith("C"):
            return True
        return False

    def _is_whitespace(self, char):
        if char == " " or char == "\t" or char == "\n" or char == "\r":
            return True
        cat = unicodedata.category(char)
        if cat == "Zs":
            return True
        return False

    def _is_punctuation(self, char):
        """Checks whether `chars` is a punctuation character."""
        cp = ord(char)
        if ((cp >= 33 and cp <= 47) or (cp >= 58 and cp <= 64) or
                (cp >= 91 and cp <= 96) or (cp >= 123 and cp <= 126)):
            return True
        cat = unicodedata.category(char)
        if cat.startswith("P"):
            return True
        return False

    def convert_to_unicode(self, text):
        """Converts `text` to Unicode (if it's not already), assuming utf-8 input."""
        if isinstance(text, str):
            return text
        elif isinstance(text, bytes):
            return text.decode("utf-8", "ignore")
        else:
            raise ValueError("Unsupported string type: %s" % (type(text)))

    def _is_chinese_char(self, cp):
        """Checks whether CP is the codepoint of a CJK character."""
        if ((cp >= 0x4E00 and cp <= 0x9FFF) or  #
                (cp >= 0x3400 and cp <= 0x4DBF) or  #
                (cp >= 0x20000 and cp <= 0x2A6DF) or  #
                (cp >= 0x2A700 and cp <= 0x2B73F) or  #
                (cp >= 0x2B740 and cp <= 0x2B81F) or  #
                (cp >= 0x2B820 and cp <= 0x2CEAF) or
                (cp >= 0xF900 and cp <= 0xFAFF) or  #
                (cp >= 0x2F800 and cp <= 0x2FA1F)):  #
            return True

        return False

    def lowercase_and_normalize(self, text):
        """转小写，并进行简单的标准化
        """
        text = text.lower()
        text = unicodedata.normalize('NFD', text)
        text = ''.join([ch for ch in text if unicodedata.category(ch) != 'Mn'])
        return text

    
    def _is_special(self, ch):
        """判断是不是有特殊含义的符号
        """
        return bool(ch) and (ch[0] == '[') and (ch[-1] == ']')

    def rematch(self, text, tokens):
        """给出原始的text和tokenize后的tokens的映射关系
        """

        if self.do_lower_case:
            text = text.lower()

        normalized_text, char_mapping = '', []
        for i, ch in enumerate(text):
            if self.do_lower_case:
                ch = self.lowercase_and_normalize(ch)
            ch = ''.join([
                c for c in ch
                if not (ord(c) == 0 or ord(c) == 0xfffd or self._is_control(c))
            ])
            normalized_text += ch
            char_mapping.extend([i] * len(ch))

        text, token_mapping, offset = normalized_text, [], 0
        for token in tokens:
            if self._is_special(token):
                token_mapping.append([])
            else:
                start = text[offset:].index(token) + offset
                end = start + len(token)
                # token_mapping.append(char_mapping[start:end])
                token_mapping.append((start, end))
                offset = end

        return token_mapping

    
    def tokenize(self, text) -> List:
        origin_text = text
        text = self._clean_text(text=origin_text)
        text = self._tokenize_chinese_chars(text)
        orig_tokens = self._whitespace_tokenize(text)
        split_tokens = []
        for token in orig_tokens:
            if self.do_lower_case and token not in self.never_split:
                token = token.lower()
                token = self._run_strip_accents(token)
            split_tokens.extend(self._run_split_on_punc(token))

        output_tokens = self._whitespace_tokenize(" ".join(split_tokens))
        if self.return_offset_mapping:
            token_mapping = self.rematch(origin_text, output_tokens)
            return output_tokens, token_mapping
        else:
            return output_tokens
        
        

@DefaultV1Recipe.register(
    DefaultV1Recipe.ComponentType.MESSAGE_TOKENIZER, is_trainable=False)
class ChineseTokenizer(Tokenizer):

    @staticmethod
    def required_packages() -> List[Text]:
        """Any extra python dependencies required for this component to run."""
        return []

    @staticmethod
    def get_default_config() -> Dict[Text, Any]:
        """Returns the component's default config."""
        return {
            # This *must* be added due to the parent class.
            "intent_tokenization_flag": False,
            # This *must* be added due to the parent class.
            "intent_split_symbol": "_",
            "never_split": [],
            "case_sensitive": True,
            "return_offset_mapping": True,
            "do_lower_case": True
        }
    
    def __init__(self, config: Dict[Text, Any]) -> None:
        """Initialize the tokenizer."""
        config = {**self.get_default_config(), **config}
        super().__init__(config)
        self.case_sensitive = config["case_sensitive"]
        self.return_offset_mapping = config['return_offset_mapping']
        self.do_lower_case = config['do_lower_case']
        self.never_split = config['never_split']
        self._tokenizer = BaseTokenizer(**config)


    

    @classmethod
    def create(
        cls,
        config: Dict[Text, Any],
        model_storage: ModelStorage,
        resource: Resource,
        execution_context: ExecutionContext,
    ) -> ChineseTokenizer:
        return cls(config)


    def tokenize(self, message: Message, attribute: Text) -> List[Token]:
        orig_text = message.get(attribute)
        output_tokens, token_mapping = self._tokenizer.tokenize(orig_text)
        return [Token(text=token_text, start=position[0], end=position[1])
                for token_text, position in zip(output_tokens, token_mapping)]


@DefaultV1Recipe.register(
    DefaultV1Recipe.ComponentType.MESSAGE_TOKENIZER, is_trainable=False
)
class BlankSpacyTokenizer(Tokenizer):
    """Creates features for entity extraction."""

    @staticmethod
    def required_packages() -> List[Text]:
        """Any extra python dependencies required for this component to run."""
        return ["spacy"]

    @staticmethod
    def get_default_config() -> Dict[Text, Any]:
        """Returns the component's default config."""
        return {
            # This *must* be added due to the parent class.
            "intent_tokenization_flag": False,
            # This *must* be added due to the parent class.
            "intent_split_symbol": "_",
            # This is the spaCy language setting.
            "lang": "zh",
        }

    def __init__(self, config: Dict[Text, Any]) -> None:
        """Initialize the tokenizer."""
        config = {**self.get_default_config(), "lang": config["lang"]}
        super().__init__(config)
        self.nlp = spacy.blank(name=config["lang"])

    @classmethod
    def create(
        cls,
        config: Dict[Text, Any],
        model_storage: ModelStorage,
        resource: Resource,
        execution_context: ExecutionContext,
    ) -> BlankSpacyTokenizer:
        return cls(config)

    def tokenize(self, message: Message, attribute: Text) -> List[Token]:
        """Tokenizes the text of the provided message."""
        text = message.get(attribute)

        doc = self.nlp(text)
        tokens = [
            Token(
                text=t.text,
                start=t.idx,
            )
            for t in doc
            if t.text and t.text.strip()
        ]
        return self._apply_token_pattern(tokens)

    @classmethod
    def validate_config(cls, config: Dict[Text, Any]) -> None:
        """Validates that the component is configured properly."""
        if "lang" not in config.keys():
            raise ValueError(
                "BlankSpacyTokenizer needs language configured via `lang`."
            )

# if __name__ == "__main__":
#     config = {}
#     t = ChineseTokenizer(config=config)
#     print(t.tokenize({'text':'hi 你好么1992'}, 'text'))
    
